import chromadb.utils.embedding_functions as embedding_functions
import chromadb
import json

ollama_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key="sk-lOnR4erWFrXwpBqW24E293122a0347AdA983D7E99e0512B1",
                api_base="http://60.204.226.75:32100/v1",
                model_name="shaw/dmeta-embedding-zh"
            )
# 打开并加载JSON文件
def read_json_data_from_jsonfile(filename):
    with open(filename, 'r', encoding='utf-8') as file:
        sql_statements = json.load(file)
    return sql_statements

documents = []
# json_objects = read_json_data_from_jsonfile("../02导入训练sql数据/data/smartpole_ddl_2.json")
# for index, json_object in enumerate(json_objects, start=1):
#     ddl = json_object.get('ddl')
#     documents.append(ddl)

chroma_client = chromadb.PersistentClient(path="./db")
collection = chroma_client.get_or_create_collection(name="my_demo02_collection", embedding_function=ollama_ef)

for i, d in enumerate(documents):
  embedding = ollama_ef(d)
  collection.upsert(
    ids=[str(i)],
    embeddings=embedding,
    documents=[d]
  )
query_results = collection.query(query_texts=["查询单灯方案名称是舜华路的方案下发详情信息"], n_results=10)
#query_results = collection.query(query_embeddings=ollama_ef("环境监测设备上报的实施数据"))

print(f"查询结果：{query_results}")